from torch.utils.data import Dataset
import numpy as np

# 自定义数据集类
class SinWaveDataset(Dataset):
    def __init__(self, num_samples, sequence_length):
        self.num_samples = num_samples
        self.sequence_length = sequence_length
        self.data, self.labels = self.generate_data()

    def generate_data(self):
        data = []
        labels = []

        for _ in range(self.num_samples):
            t = np.linspace(0, 4*np.pi, self.sequence_length)  # 时间点
            freq = np.random.uniform(1.0, 10.0)  # 随机频率
            amplitude = np.random.uniform(0.5, 2.0)  # 随机振幅
            phase = np.random.uniform(0, 2 * np.pi)  # 随机相位
            b = np.random.uniform(-2.0, 2.0)  # 随机偏置
            x = amplitude * np.sin(freq * t + phase) + b   # 生成x数值
            condition=np.array([amplitude, freq, phase])
            data.append(x)  # 每个样本是 (sequence_length)
            labels.append(condition)  # 振幅、频率、相位作为标签

        data = np.array(data, dtype=np.float32)
        labels = np.array(labels, dtype=np.float32)
        return data, labels

    def __len__(self):
        return self.num_samples

    def __getitem__(self, idx):
        return self.data[idx], self.labels[idx]